#include "..\include\socket.hpp"

#define WIN32_LEAN_AND_MEAN

#include <windows.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <string>
#include <sstream>

// Need to link with Ws2_32.lib
#pragma comment(lib, "Ws2_32.lib")

using namespace std;

#define ERR(fun, code) ErrorMessage(__FILE__, __LINE__, __FUNCTION__, fun, code)
#define E(fun)                        \
    do                                \
    {                                 \
        DWORD dw = WSAGetLastError(); \
        if (dw)                       \
            throw ERR(fun, dw);       \
    } while (false)

string ErrorMessage(LPCSTR lpszFileName, DWORD dwLine, LPCSTR lpszParentFunction, LPCSTR lpszFunction, DWORD dwCode); // 输出错误信息

WSA::WSA()
{
    WSADATA wsaData;
    // Initialize Winsock
    int r = WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (r)
        throw ERR("WSAStartup", r);
}

WSA::~WSA()
{
    WSACleanup();
}

struct socket_internal_t
{
    SOCKET sock = INVALID_SOCKET;
    ~socket_internal_t()
    {
        closesocket(sock);
    }
    void find_available_socket(const char *host, int port, addrinfo &hints, int(WINAPI *fun)(SOCKET, const sockaddr *, int));
    int find_available_socket1(const char *host, int port, addrinfo &hints, int(WINAPI *fun)(SOCKET, const char *, int, int, const struct sockaddr*, int), const char* buf, int len, int flag);
};

socket_t::socket_t() : ptr(new socket_internal_t)
{
}

socket_t::~socket_t()
{
    delete ptr;
}

void init_addrinfo(addrinfo &hints, SOCK_TYPE type)
{
    switch (type)
    {
    case SOCK_TYPE::RAW:
        hints.ai_socktype = SOCK_RAW;
        hints.ai_protocol = IPPROTO_RAW;
        break;
    case SOCK_TYPE::TCP:
        hints.ai_socktype = SOCK_STREAM;
        hints.ai_protocol = IPPROTO_TCP;
        break;
    case SOCK_TYPE::UDP:
        hints.ai_socktype = SOCK_DGRAM;
        hints.ai_protocol = IPPROTO_UDP;
        break;
    }
}

void client_t::connect(const char *host, int port, SOCK_TYPE type)
{
    addrinfo hints = {0};
    hints.ai_family = AF_UNSPEC;
    init_addrinfo(hints, type);

    reinterpret_cast<socket_internal_t *>(ptr)->find_available_socket(host, port, hints, ::connect);
}

void server_t::bind(const char *host, int port, SOCK_TYPE type)
{
    addrinfo hints = {0};
    hints.ai_family = AF_INET;
    hints.ai_flags = AI_PASSIVE;
    init_addrinfo(hints, type);

    reinterpret_cast<socket_internal_t *>(ptr)->find_available_socket(host, port, hints, ::bind);
}

void server_t::listen(int backlog)
{
    ::listen(reinterpret_cast<socket_internal_t *>(ptr)->sock, backlog);
}

void server_t::listen()
{
    listen(SOMAXCONN);
}

void server_t::accept()
{
    auto &sock = reinterpret_cast<socket_internal_t *>(ptr)->sock;
    SOCKET s = ::accept(sock, NULL, NULL);
    if (s == INVALID_SOCKET)
        E("accept");
    closesocket(sock);
    sock = s;
}

void socket_t::shutdown(bool read, bool write)
{
    int how;
    if (read && write)
        how = SD_BOTH;
    else if (read)
        how = SD_RECEIVE;
    else if (write)
        how = SD_SEND;
    else
        return;
    ::shutdown(reinterpret_cast<socket_internal_t *>(ptr)->sock, how);
}

int socket_t::send(const string &buf)
{
    int r = ::send(reinterpret_cast<socket_internal_t *>(ptr)->sock, buf.data(), buf.size(), 0);
    if (r == SOCKET_ERROR)
        E("send");
    return r;
}

string socket_t::recv()
{
    string s;
    int r;
    do
    {
        char buf[8];
        r = ::recv(reinterpret_cast<socket_internal_t *>(ptr)->sock, buf, sizeof buf, 0);
        if (r > 0)
            s.append(buf, buf + r);
        else if (r < 0)
            E("recv");
    } while (r > 0);
    return s;
}

int socket_t::sendto(const std::string &buf, const char *host, int port, SOCK_TYPE type)
{
    addrinfo hints = {0};
    hints.ai_family = AF_UNSPEC;
    init_addrinfo(hints, type);

    return reinterpret_cast<socket_internal_t *>(ptr)->find_available_socket1(host, port, &hints, ::sendto, buf.data(), (int)buf.size(), 0);
}

recv_data socket_t::recvfrom()
{
    recv_data rd;
    char buf[0x1000];
    struct sockaddr_in addr_in;
    socklen_t addr_len;
    int r = ::recvfrom(reinterpret_cast<socket_internal_t *>(ptr)->sock, buf, sizeof(buf), 0, (struct sockaddr *)&addr_in, &addr_len);
    if (r > 0)
        rd.data.append(buf, buf + r);
    else if (r < 0)
        E("recv");
    rd.host = inet_ntoa(addr_in.sin_addr);
    rd.port = addr_in.sin_port;
    return rd;
}

void socket_internal_t::find_available_socket(const char *host, int port, addrinfo &hints, int(WINAPI *fun)(SOCKET, const sockaddr *, int))
{
    class addrinfo_t
    {
    public:
        addrinfo_t() : d(nullptr) {}
        ~addrinfo_t()
        {
            if (d)
                freeaddrinfo(d);
        }
        addrinfo **operator&()
        {
            return &d;
        }
        operator addrinfo *()
        {
            return d;
        }

    private:
        addrinfo *d;
    } result;

    // Resolve the server address and port
    ostringstream oss;
    oss << port;
    int r = getaddrinfo(host, oss.str().c_str(), &hints, &result);
    if (r)
        throw ERR("getaddrinfo", r);

    // Attempt to connect to an address until one succeeds
    for (addrinfo *ptr = result; ptr; ptr = ptr->ai_next)
    {
        // Create a SOCKET for connecting to server
        sock = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
        if (sock == INVALID_SOCKET)
            E("socket");

        r = fun(sock, ptr->ai_addr, (int)ptr->ai_addrlen);
        if (r != SOCKET_ERROR)
            return;
        sock = INVALID_SOCKET;
    }

    using namespace string_literals;

    if (sock == INVALID_SOCKET)
    {
        if (hints.ai_flags)
            throw "Unable to bind to port!"s;
        else
            throw "Unable to connect to server!"s;
    }
}

int socket_internal_t::find_available_socket(const char *host, int port, addrinfo &hints, int(WINAPI *fun)(SOCKET, const char *, int, int, const struct sockaddr*, int), const char* buf, int len, int flag)
{
    class addrinfo_t
    {
    public:
        addrinfo_t() : d(nullptr) {}
        ~addrinfo_t()
        {
            if (d)
                freeaddrinfo(d);
        }
        addrinfo **operator&()
        {
            return &d;
        }
        operator addrinfo *()
        {
            return d;
        }

    private:
        addrinfo *d;
    } result;

    // Resolve the server address and port
    ostringstream oss;
    oss << port;
    int r = getaddrinfo(host, oss.str().c_str(), &hints, &result);
    if (r)
        throw ERR("getaddrinfo", r);

    // Attempt to connect to an address until one succeeds
    for (addrinfo *ptr = result; ptr; ptr = ptr->ai_next)
    {
        // Create a SOCKET for connecting to server
        sock = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
        if (sock == INVALID_SOCKET)
            E("socket");

        r = fun(sock, buf, len, flag, ptr->ai_addr, (int)ptr->ai_addrlen);
        if (r != SOCKET_ERROR)
            return r;
        sock = INVALID_SOCKET;
    }

    using namespace string_literals;

    if (sock == INVALID_SOCKET)
    {
        if (hints.ai_flags)
            throw "Unable to bind to port!"s;
        else
            throw "Unable to connect to server!"s;
    }
    return 0;
}

string ErrorMessage(LPCSTR lpszFileName, DWORD dwLine, LPCSTR lpszParentFunction, LPCSTR lpszFunction, DWORD dwCode)
{
    // Retrieve the system error message for the last-error code
    LPSTR lpMsgBuf = NULL;
    LPSTR lpDisplayBuf = NULL;

    FormatMessageA(
        FORMAT_MESSAGE_ALLOCATE_BUFFER |
            FORMAT_MESSAGE_FROM_SYSTEM |
            FORMAT_MESSAGE_IGNORE_INSERTS |
            FORMAT_MESSAGE_MAX_WIDTH_MASK,
        NULL,
        dwCode,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPSTR)&lpMsgBuf,
        0, NULL);

    ostringstream oss;
    oss << "Filename: " << lpszFileName << endl
        << "Line: " << dwLine << endl
        << "Parent function: " << lpszParentFunction << endl
        << "Function: " << lpszFunction << endl
        << "Code: " << dwCode;

    if (lpMsgBuf)
    {
        oss << endl
            << "Description: " << lpMsgBuf;
        LocalFree(lpMsgBuf);
    }

    return oss.str();
}